import matplotlib.pyplot as plt
import numpy as np
from sklearn.neighbors import KNeighborsClassifier
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split

#return_X_y = True为了返回的数据是元组而不是字典
X, y = load_iris(return_X_y = True)
#将数据集进行划分，训练集占7份，测试集占3份
train_X, test_X, train_y, test_y = train_test_split(X, y, test_size = 0.3, train_size = 0.7, stratify = y, random_state = 42)
scores = []
#使用不同的邻居数进行训练测试
for n in range(1, 6) :
    knn = KNeighborsClassifier(n_neighbors = n)
    #训练
    knn.fit(train_X, train_y)
    #预测
    pred = knn.predict(test_X)
    #准确率并保留3位小数
    score = round(knn.score(test_X, test_y), 3)
    scores.append(score)

print(scores)
plt.figure()
plt.plot(range(1, 6), scores, 'o--', color = 'blue')
plt.xlabel('$n\_neighbors$', fontsize = 14)
plt.ylabel('$precision$', fontsize = 14)
for x, y in zip(range(1, 6), scores) :
    plt.text(x - 0.18, y - 0.1, f'${y}$', fontsize = 14)
plt.title(f'$precision\ of\ different\ neighors$', fontsize = 14)
plt.xticks(np.arange(1, 6))
plt.yticks(np.linspace(0, 1, 5))
plt.show()
plt.savefig('figure.png')
